# -*- coding: utf-8 -*-

# File Name： resnet
# Description :
# Author : lirui
# create_date： 2022/6/3
# Change Activity:
from torch import nn
from torch.nn import Conv2d

from ..base.base_module import BaseModule
from ..layers.norm import build_norm_layer
from ..modules.res_layer import ResLayer
from ..modules.resnet_block import ResnetBlock
from ...misc.util import int_to_tuple


class ResNet(BaseModule):
    arch_settings = {
        18: (ResnetBlock, (2, 2, 2, 2)),
        34: (ResnetBlock, (3, 4, 6, 3)),
        50: (ResnetBlock, (3, 4, 6, 3)),  # Bottleneck
        101: (ResnetBlock, (3, 4, 23, 3)),  # Bottleneck
        152: (ResnetBlock, (3, 8, 36, 3))  # Bottleneck
    }

    def __init__(self,
                 depth=18,
                 in_channels=3,
                 stem_channels=64,
                 deep_stem=False,
                 strides=(1, 2, 2, 2),
                 dilations=(1, 1, 1, 1),
                 out_indices=(0, 1, 2, 3),
                 expansion=1,
                 layers_channels=(256, 512, 1024, 2048),
                 norm_cfg=None
                 ):
        super().__init__()
        if norm_cfg is None:
            norm_cfg = dict(type='BN')
        self.in_channels = in_channels
        self.stem_channels = stem_channels
        self.block, self.layers_num = self.arch_settings[depth]
        self.norm_cfg = norm_cfg
        self.deep_stem = deep_stem
        self.strides = strides
        self.dilations = dilations
        self.expansion = expansion
        self.out_indices = out_indices
        self.layers_channels = layers_channels
        self.stem, self.maxpool = self._build_stem(in_channels, stem_channels)
        self.layers = self._build_layers()
        pass

    def _build_layers(self):
        """
        build  backbone layers.

        Returns:

        """
        layers = []
        in_channel = self.stem_channels
        for i, num_layer in enumerate(self.layers_num):
            stride = self.strides[i]
            dilation = self.dilations[i]
            channel = self.layers_channels[i]
            res_layer = ResLayer(
                block=self.block,
                in_channels=in_channel,
                out_channels=channel,
                num_blocks=num_layer,
                stride=stride,
                dilation=dilation,
            )
            in_channel = channel * self.expansion
            layer_name = f'layer{i + 1}'
            self.add_module(layer_name, res_layer)
            layers.append(res_layer)
        return layers

    def forward(self, x):
        """Forward function."""
        x = self.stem(x)
        x = self.maxpool(x)
        outs = []
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i in self.out_indices:
                outs.append(x)
        return tuple(outs)

    def _build_stem(self, in_channels, stem_channels):
        """

        Args:
            in_channels:
            stem_channels:

        Returns:

        """
        middle_norm_cfg = self.norm_cfg.copy()
        middle_norm_cfg['num_features'] = stem_channels // 2
        output_norm_cfg = self.norm_cfg.copy()
        output_norm_cfg['num_features'] = stem_channels
        if self.deep_stem:
            stem = nn.Sequential(
                Conv2d(in_channels,
                       stem_channels // 2,
                       kernel_size=int_to_tuple(3),
                       stride=int_to_tuple(2),
                       padding=1,
                       bias=False),
                build_norm_layer(middle_norm_cfg),
                nn.ReLU(inplace=True),
                Conv2d(stem_channels // 2,
                       stem_channels // 2,
                       kernel_size=int_to_tuple(3),
                       stride=int_to_tuple(1),
                       padding=1,
                       bias=False),
                build_norm_layer(middle_norm_cfg),
                nn.ReLU(inplace=True),
                Conv2d(stem_channels // 2,
                       stem_channels,
                       kernel_size=int_to_tuple(3),
                       stride=int_to_tuple(1),
                       padding=1,
                       bias=False),
                build_norm_layer(output_norm_cfg),
                nn.ReLU(inplace=True))
        else:
            stem = nn.Sequential(
                Conv2d(in_channels,
                       out_channels=self.stem_channels,
                       kernel_size=int_to_tuple(7),
                       stride=int_to_tuple(2),
                       padding=3,
                       bias=False),
                build_norm_layer(output_norm_cfg),
                nn.ReLU(inplace=True)
            )
        maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        return stem, maxpool
